In `@src/mini_trainer/train.py`:
- Around line 1235-1238: The MLflow enablement check currently only sets
use_mlflow based on mlflow_tracking_uri; change it so use_mlflow is true if any
MLflow-related arg is provided (e.g., mlflow_tracking_uri,
mlflow_experiment_name, mlflow_run_name) so passing --mlflow-experiment-name or
--mlflow-run-name enables MLflow even when tracking URI is omitted; update the
initialization of use_mlflow (the variable referenced in this file) to check all
MLflow args and keep use_wandb logic unchanged.
- Around line 1304-1313: The MLflow init and log_params block (inside the
use_mlflow branch calling mlflow_wrapper.init and mlflow_wrapper.log_params,
currently gated by local_rank == 0 via log_rank_0) must be executed only on the
global rank 0 process to avoid creating one MLflow run per node; change the
guard so that the mlflow_wrapper.init and mlflow_wrapper.log_params calls run
only when the global process rank is 0 (e.g., use torch.distributed.get_rank()
== 0 or an existing global_rank variable) and keep
AsyncStructuredLogger/log_rank_0 behavior unchanged for other logs.

In `@src/mini_trainer/mlflow_wrapper.py`:
- Around line 33-78: log_params() and log() re-enter an already active MLflow
run using mlflow.start_run(run_id=_active_run_id) which raises unless
nested=True is set; update both functions (log_params and log) so their with
mlflow.start_run(...) calls include nested=True (i.e., with
mlflow.start_run(run_id=_active_run_id, nested=True):) to allow re-entrancy into
the existing run when _active_run_id is set.
- Around line 81-128: log_params and log unconditionally call
mlflow.start_run(run_id=_active_run_id) which raises if a run is already active;
modify both functions (log_params, log) to first call mlflow.active_run() and
only invoke mlflow.start_run(run_id=_active_run_id) when active_run is None or
active_run.info.run_id != _active_run_id, otherwise skip starting a run and
directly call mlflow.log_params / mlflow.log_metrics; reference _active_run_id,
mlflow.active_run(), mlflow.start_run(), mlflow.log_params, and
mlflow.log_metrics to locate where to add the conditional guard.

In `@src/mini_trainer/mlflow_wrapper.py`:
- Around line 96-110: The log_params function currently assumes an active mlflow
run; update mlflow_wrapper.log_params to mirror the async-safe pattern used in
log(): after calling check_mlflow_available("log params to mlflow"), detect
mlflow.active_run() and if it's None use the module-level fallback
_active_run_id to call mlflow.start_run(run_id=_active_run_id) (or use
mlflow.log_params with the run context started) so parameters are logged against
the correct run when thread-local context is lost; ensure you still convert
params to strings before logging and maintain the same error behavior.

In `@src/mini_trainer/mlflow_wrapper.py`:
- Around line 110-119: The current logging uses
mlflow.start_run(run_id=_active_run_id) which as a context manager will call
mlflow.end_run() and finish the run; instead, when active_run is None but
_active_run_id is set, use MlflowClient to log directly by run id (e.g., replace
the with mlflow.start_run(...) blocks around mlflow.log_params and the similar
block for metrics) by calling MlflowClient().log_params(run_id=_active_run_id,
params=str_params) and MlflowClient().log_metric(s)/log_batch for metrics, so
the stored run (_active_run_id) is not ended prematurely and future logs
continue to target the same run.

In `@src/mini_trainer/mlflow_wrapper.py`:
- Around line 45-90: The init function can raise TypeError if callers include
run_name in kwargs as well as the explicit run_name parameter; before calling
mlflow.start_run, pop and prefer the explicit run_name from kwargs (e.g., if
"run_name" in kwargs: kwargs.pop("run_name")) so the explicit run_name argument
takes precedence and no duplicate keyword is passed to mlflow.start_run; update
the call site (mlflow.start_run(run_name=run_name, **kwargs)) in init to use the
sanitized kwargs.

In `@src/mini_trainer/api_train.py`:
- Around line 157-165: The MLflow experiment/run name flags are currently only
added when train_args.mlflow_tracking_uri is set; update the logic in
src/mini_trainer/api_train.py so that the appending of --mlflow-experiment-name
and --mlflow-run-name is done independently of the --mlflow-tracking-uri block:
keep the existing append of
f"--mlflow-tracking-uri={train_args.mlflow_tracking_uri}" inside the if
train_args.mlflow_tracking_uri: block, but move the if
train_args.mlflow_experiment_name: and if train_args.mlflow_run_name: checks out
of that block so they append their flags to the command list even when
mlflow_tracking_uri is None (consistent with use_mlflow = any([...])).

In `@src/mini_trainer/mlflow_wrapper.py`:
- Around line 28-42: The error message in check_mlflow_available currently
suggests installing "mini-trainer[mlflow]" which doesn't match the package name;
update the error_msg string used in check_mlflow_available (and raised via
MLflowNotAvailableError) to instruct installing
"rhai-innovation-mini-trainer[mlflow]" instead, keeping the same phrasing and
logger.error call and ensuring MLFLOW_AVAILABLE and MLflowNotAvailableError
usage remain unchanged.
- Around line 45-99: The current init function unconditionally does
kwargs.pop("run_name", None), which discards a caller-supplied run name passed
via **kwargs; change the logic so that if the explicit parameter run_name is
None and "run_name" exists in kwargs, promote it by assigning run_name =
kwargs.pop("run_name"), otherwise (when run_name is explicitly provided) remove
"run_name" from kwargs to avoid duplicate arguments before calling
mlflow.start_run; update references around init, run_name, kwargs.pop and
mlflow.start_run accordingly.